from src.dataset import DeepFashionCAPDataset
from src.const import base_path
from src.networks import VGG16BaselineNet
import pandas as pd
import torch
import torch.utils.data
import numpy as np
from src import const
from tensorboardX import SummaryWriter
import time
import os
from watchdog.observers import Observer
import watchdog
from utils import Evaluator

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
writer = SummaryWriter()


class Handler(watchdog.events.FileSystemEventHandler):

    def __init__(self, evaluate_path, net, val_dataloader):
        super(Handler, self).__init__()
        self.is_evaluating = False
        self.net = net
        self.val_dataloader = val_dataloader
        self.evaluate_path = evaluate_path

    def on_modified(self, event):
        if event.src_path.endswith(evaluate_path):
            if self.is_evaluating:
                pass
            else:
                self.is_evaluating = True
                time.sleep(60)
                evaluate(self.net, event.src_path, self.val_dataloader)
                self.is_evaluating = False


def evaluate(net, net_path, val_dataloader):
    val_step = len(val_dataloader)
    net.load_state_dict(torch.load(net_path))
    net.eval()
    evaluator = Evaluator(const.VAL_CATEGORY_TOP_N, const.VAL_LM_RELATIVE_DIS)
    with torch.no_grad():
        if getattr(net, 'step'):
            step = net.step
        else:
            step = int(time.time())  # use timestamp if there is no step information
        for i, sample in enumerate(val_dataloader):
            for key in sample:
                sample[key] = sample[key].to(device)
            output = net(sample)
            evaluator.add(output, sample)
            if (i + 1) % 100 == 0:
                print('Val Step [{}/{}]'.format(i + 1, val_step))
        ret = evaluator.evaluate()
        for topk, accuracy in ret['category_accuracy_topk'].items():
            writer.add_scalar('metrics/category_top{}'.format(topk), accuracy, step)
            print('category top', topk, 'accuracy:', accuracy)
        writer.add_scalar('metrics/attr_recall', ret['attr_recall'], step)
        print('attr_recall:', ret['attr_recall'])
        writer.add_scalar('metrics/attr_precision', ret['attr_precision'], step)
        print('attr_precision', ret['attr_precision'])
        writer.add_scalar('metrics/lm_vis_precision', ret['lm_vis_precision'], step)
        print('lm_vis_precision', ret['lm_vis_precision'])
        writer.add_scalar('metrics/lm_pos_precision', ret['lm_pos_precision'], step)
        print('lm_pos_precision', ret['lm_pos_precision'])



if __name__ == '__main__':
    evaluate_path = 'models/vgg16.pkl'

    df = pd.read_csv(base_path + 'info.csv')
    val_df = df[df['evaluation_status'] == 'val']
    val_dataset = DeepFashionCAPDataset(val_df, mode='CENTER')
    val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=const.BATCH_SIZE * 4, shuffle=True, num_workers=4)
    net = VGG16BaselineNet().to(device)
    if os.path.exists(evaluate_path):
        evaluate(net, evaluate_path, val_dataloader)
    observer = Observer()
    observer.schedule(Handler(evaluate_path, net, val_dataloader), 'models/', recursive=False)
    print('Start Observation~')
    observer.start()
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        observer.stop()
    observer.join()
